﻿using System;
using System.Collections.Generic;

namespace Novice.Tools
{
    public enum TreeEnumerationStrategy
    {
        RootFirst,
        RootLast
    }

    public static class EnumTree
    {
        public static IEnumerable<T> FlattenSingleItem<T>(
            T item, 
            Func<T, IEnumerator<T>> traceFunction, 
            TreeEnumerationStrategy enumerationStrategy)
        {
            return new FlattenedTreeCollection<T>(
                new SingleItemCollection<T>(item),
                traceFunction,
                GetTreeEnumerationStrategy<T>(enumerationStrategy));
        }

        public static IEnumerable<T> FlattenSingleItem<T>(
            T item, 
            Func<T, IEnumerable<T>> traceFunction,
            TreeEnumerationStrategy enumerationStrategy)
        {
            return FlattenSingleItem(item,
                                     a =>
                                     {
                                         var c = traceFunction(a);
                                         return (c == null) ? null : c.GetEnumerator();
                                     },
                                     enumerationStrategy);
        }

        public static IEnumerable<T> Flatten<T>(
            this IEnumerable<T> collection, 
            Func<T, IEnumerator<T>> traceFunction, 
            TreeEnumerationStrategy enumertionStrategy = TreeEnumerationStrategy.RootFirst)
        {
            return new FlattenedTreeCollection<T>(
                collection,
                traceFunction,
                GetTreeEnumerationStrategy<T>(enumertionStrategy));
        }

        public static IEnumerable<T> Flatten<T>(
            this IEnumerable<T> collection, 
            Func<T, IEnumerable<T>> traceFunction,
            TreeEnumerationStrategy enumertionStrategy = TreeEnumerationStrategy.RootFirst)
        {
            return collection.Flatten(a =>
                                      {
                                          var c = traceFunction(a);
                                          return (c == null) ? null : c.GetEnumerator();
                                      }, 
                                      enumertionStrategy);
        }

        public static IEnumerable<T> Flatten<T>(this IEnumerable<T> collection)
            where T : IEnumerable<T>
        {
            return collection.Flatten(a => a.GetEnumerator());
        }

        private static ITreeEnumertionStrategy<T> GetTreeEnumerationStrategy<T>(TreeEnumerationStrategy strategy)
        {
            switch (strategy)
            {
                case TreeEnumerationStrategy.RootFirst:
                    return new RootFirstTreeEnumerationStrategy<T>();
                case TreeEnumerationStrategy.RootLast:
                    return new RootLastTreeEnumerationStrategy<T>();
                default:
                    throw new ApplicationException("Unknown type of tree enumeration");
            }
        }

    }


    /// <summary>
    /// IEnumerator, содержащий единственный элемент
    /// </summary>
    /// <typeparam name="T">Тип элемента</typeparam>
    internal class SingleItemEnumerator<T>: IEnumerator<T>
    {
        #region private data

        private readonly T _item;
        private bool _enumerationIsOver;

        #endregion

        public SingleItemEnumerator(T item)
        {
            _item = item;
            _enumerationIsOver = false;
        }

        public T Current
        {
            get {
                return _enumerationIsOver ? default(T) : _item;
            }
        }

        public void Dispose()
        {
            // ничего не делать
        }

        object System.Collections.IEnumerator.Current
        {
            get
            {
                return _enumerationIsOver ? default(T) : _item;
            }
        }

        public bool MoveNext()
        {
            return !_enumerationIsOver;
        }

        public void Reset()
        {
            _enumerationIsOver = false;
        }
    }


    internal struct TreeEnumerationContext<T>: IDisposable
    {
        public TreeEnumerationContext(IEnumerator<T> root, Func<T, IEnumerator<T>> traceFunction)
        {
            GetEnumerator = traceFunction;

            Stack = new Stack<IEnumerator<T>>();    
            Stack.Push(root);
        }

        public readonly Stack<IEnumerator<T>> Stack;
        public readonly Func<T, IEnumerator<T>> GetEnumerator;

        public void Dispose()
        {
            while (Stack.Count > 1)
            {
                Stack.Pop().Dispose();
            }
        }
    }

    internal interface ITreeEnumertionStrategy<T>
    {
        bool MoveNext(TreeEnumerationContext<T> context, out T nextItem);
    }
    
    /// <summary>
    /// Базовый класс для древовидных перечислителей
    /// </summary>
    /// <typeparam name="T">Тип элмента</typeparam>
    internal class TreeEnumerator<T> : IEnumerator<T>
    {
        #region protected data

        private readonly IEnumerator<T> _root;
        private readonly Func<T, IEnumerator<T>> _traceFunction;
        
        private TreeEnumerationContext<T> _context;

        private readonly ITreeEnumertionStrategy<T> _enumertionStrategy;
        private T _currentItem;

        #endregion

        public TreeEnumerator(
            IEnumerator<T> root, 
            Func<T, IEnumerator<T>> traceFunction, 
            ITreeEnumertionStrategy<T> enumertionStrategy)
        {
            _root = root;
            _traceFunction = traceFunction;
            _enumertionStrategy = enumertionStrategy;

            _context = new TreeEnumerationContext<T>(root, traceFunction);
            _currentItem = default(T);
        }

        ~TreeEnumerator()
        {
            Dispose(false);
        }

        public T Current
        {
            get { return _currentItem; }
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!disposing) return;

            _context.Dispose();
            _root.Dispose();
        }

        public void Reset()
        {
            var oldContext = _context;
            _context = new TreeEnumerationContext<T>(_root, _traceFunction);
            oldContext.Dispose();

            _root.Reset();
        }

        object System.Collections.IEnumerator.Current
        {
            get { return _currentItem; }
        }

        public bool MoveNext()
        {
            return _enumertionStrategy.MoveNext(_context, out _currentItem);
        }
    }

    internal class RootFirstTreeEnumerationStrategy<T> : ITreeEnumertionStrategy<T>
    {
        public bool MoveNext(TreeEnumerationContext<T> context, out T nextItem)
        {
            var stack = context.Stack;
            var getEnumerator = context.GetEnumerator;

            if (stack.Count == 0)
            {
                nextItem = default(T);
                return false;
            }

            while (!stack.Peek().MoveNext())
            {
                // выталкивание
                var popped = stack.Pop();

                if (stack.Count == 0)
                {
                    nextItem = default(T);
                    return false;
                }

                popped.Dispose();
            }

            // заталкивание
            var current = stack.Peek().Current;
            var newEnum = getEnumerator(current);
            if (newEnum != null)
                stack.Push(newEnum);

            // устанавливаем текущий элемент
            nextItem = current;
            return true;
        }
    }

    internal class RootLastTreeEnumerationStrategy<T> : ITreeEnumertionStrategy<T>
    {
        public bool MoveNext(TreeEnumerationContext<T> context, out T nextItem)
        {
            var stack = context.Stack;
            var getEnumerator = context.GetEnumerator;

            if (stack.Count == 0)
            {
                nextItem = default(T);
                return false;
            }

            if (stack.Peek().MoveNext())
            {
                // заталкивание
                var current = stack.Peek().Current;
                var newEnum = getEnumerator(current);
                while (newEnum != null && newEnum.MoveNext())
                {
                    stack.Push(newEnum);
                    current = stack.Peek().Current;
                    newEnum = getEnumerator(current);
                }

                // устанавливаем текущий элемент
                nextItem = current;
                return true;
            }

            // выталкивание
            var popped = stack.Pop();

            if (stack.Count == 0)
            {
                nextItem = default(T);
                return false;
            }

            popped.Dispose();
            nextItem = stack.Peek().Current;
            return true;
        }
    }
   
    internal class FlattenedTreeCollection<T>: IEnumerable<T>
    {
        #region private data

        private readonly IEnumerable<T> _collection;
        private readonly Func<T, IEnumerator<T>> _traceFunction;
        private readonly ITreeEnumertionStrategy<T> _enumerationStrategy; 

        #endregion

        public FlattenedTreeCollection(
            IEnumerable<T> collection, 
            Func<T, IEnumerator<T>> traceFunction, 
            ITreeEnumertionStrategy<T> enumertionStrategy )
        {
            _collection = collection;
            _traceFunction = traceFunction;
            _enumerationStrategy = enumertionStrategy;
        }

        public IEnumerator<T> GetEnumerator()
        {
            return new TreeEnumerator<T>(
                _collection.GetEnumerator(), 
                _traceFunction,
                _enumerationStrategy);
        }

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return new TreeEnumerator<T>(
                _collection.GetEnumerator(), 
                _traceFunction, 
                _enumerationStrategy);
        }
    }

    internal class SingleItemCollection<T>: IEnumerable<T>
    {
        #region private data

        private readonly T _item;

        #endregion

        public SingleItemCollection(T item)
        {
            _item = item;
        }

        public IEnumerator<T> GetEnumerator()
        {
            return new SingleItemEnumerator<T>(_item);
        }

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return new SingleItemEnumerator<T>(_item);
        }
    }

}

